#  ------------------------------------------------------------------------------------------
#  Copyright (c) Microsoft Corporation. All rights reserved.
#  Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
#  ------------------------------------------------------------------------------------------
import math
from typing import List

import torch
import torch.nn as nn
import torch.nn.functional as F


class LoRALayer():
    def __init__(
            self,
            r: int,  # 秩
            lora_alpha: int,
            lora_dropout: float,
            merge_weights: bool,  # 是否合并权重
    ):
        self.r = r
        self.lora_alpha = lora_alpha
        # Optional dropout
        if lora_dropout > 0.:
            self.lora_dropout = nn.Dropout(p=lora_dropout)
        else:
            # 直接返回输入，表示没有dropout
            self.lora_dropout = lambda x: x
        # 表示权重尚未合并
        self.merged = False
        # 是否合并权重
        self.merge_weights = merge_weights


class Embedding(nn.Embedding, LoRALayer):
    # LoRA implemented in a dense layer
    def __init__(
            self,
            num_embeddings: int,  # 词汇大小
            embedding_dim: int,  # 每个嵌入向量的维度
            r: int = 0,  # LORA的秩
            lora_alpha: int = 1,
            merge_weights: bool = True,  # 是否合并权重
            **kwargs
    ):
        # 初始化嵌入层
        nn.Embedding.__init__(self, num_embeddings, embedding_dim, **kwargs)
        # 初始化LORA相关参数
        LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=0,
                           merge_weights=merge_weights)
        # Actual trainable parameters
        # 初始化实际的可训练参数
        if r > 0:
            # 初始化为零矩阵，够创建一个和 self.weight 数据类型、设备（CPU/GPU）、dtype（数据精度）一致的全零张量，但形状为 (r, num_embeddings)
            self.lora_A = nn.Parameter(self.weight.new_zeros((r, num_embeddings)))
            self.lora_B = nn.Parameter(self.weight.new_zeros((embedding_dim, r)))
            # 缩放因子,用于缩放LORA参数
            self.scaling = self.lora_alpha / self.r
            # Freezing the pre-trained weight matrix
            # 将预训练的权重矩阵self.weight冻结
            self.weight.requires_grad = False
        # 重置嵌入层的参数
        self.reset_parameters()

    def reset_parameters(self):
        """重置初始化参数"""
        nn.Embedding.reset_parameters(self)
        if hasattr(self, 'lora_A'):
            # initialize A the same way as the default for nn.Linear and B to zero
            # 将self.lora_A初始化为0矩阵
            nn.init.zeros_(self.lora_A)
            # 将self.lora_B初始化为服从正态分布的随机矩阵
            nn.init.normal_(self.lora_B)

    def train(self, mode: bool = True):
        """在训练模式与评估模式之间切换时处理权重合并的逻辑"""
        # 设置模块的训练模式
        nn.Embedding.train(self, mode)
        if mode:  # 训练模式
            if self.merge_weights and self.merged:
                # Make sure that the weights are not merged
                # 如果是，确保权重未被合并
                if self.r > 0:
                    # 撤销合并的权重
                    self.weight.data -= (self.lora_B @ self.lora_A).transpose(0, 1) * self.scaling
                self.merged = False
        else:  # 评估模式
            # 如果是,合并权重标记
            if self.merge_weights and not self.merged:
                # Merge the weights and mark it
                if self.r > 0:
                    # 合并权重
                    self.weight.data += (self.lora_B @ self.lora_A).transpose(0, 1) * self.scaling
                self.merged = True

    def forward(self, x: torch.Tensor):
        """利用LORA方法在嵌入层中进行低秩适配操作"""
        # 如果r>0且权重没有合并，则执行LORA适配的操作
        if self.r > 0 and not self.merged:
            # 对输入x执行嵌入操作，并将结果存储在result中
            result = nn.Embedding.forward(self, x)
            # 将输入x转换为对应的向量
            after_A = F.embedding(
                x,
                self.lora_A.transpose(0, 1),
                self.padding_idx,
                self.max_norm,
                self.norm_type,
                self.scale_grad_by_freq,
                self.sparse
            )
            result += (after_A @ self.lora_B.transpose(0, 1)) * self.scaling
            return result
        else:
            # 如果r不大于0或权重已经合并，直接调用nn.Embedding的forward方法，返回其结果
            return nn.Embedding.forward(self, x)


class Linear(nn.Linear, LoRALayer):
    # 主要实现了LORA方法在一个全连接层中的应用
    def __init__(
            self,
            in_features: int,  # 全连接层的输入和输出
            out_features: int,
            r: int = 0,  # 低秩矩阵的秩
            lora_alpha: int = 1,  # 缩放因子
            lora_dropout: float = 0.,  # dropout比例
            fan_in_fan_out: bool = False,  # 如果要替换的层存储权重的方式是（fan_in, fan_out），请将此设置为True
            merge_weights: bool = True,  # 是否合并权重
            **kwargs
    ):
        nn.Linear.__init__(self, in_features, out_features, **kwargs)
        LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout,
                           merge_weights=merge_weights)

        self.fan_in_fan_out = fan_in_fan_out
        # Actual trainable parameters
        if r > 0:
            self.lora_A = nn.Parameter(self.weight.new_zeros((r, in_features)))
            self.lora_B = nn.Parameter(self.weight.new_zeros((out_features, r)))
            self.scaling = self.lora_alpha / self.r
            # Freezing the pre-trained weight matrix
            self.weight.requires_grad = False
        self.reset_parameters()
        if fan_in_fan_out:
            self.weight.data = self.weight.data.transpose(0, 1)

    def reset_parameters(self):
        nn.Linear.reset_parameters(self)
        if hasattr(self, 'lora_A'):
            # initialize B the same way as the default for nn.Linear and A to zero
            # this is different than what is described in the paper but should not affect performance
            # 使用kaiming_uniform_方法初始化,这是一种适用于Relu激活函数的权重初始化方法
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B)

    def train(self, mode: bool = True):
        def T(w):
            # 根据fan_in_fan_out的值，决定是否对权重矩阵进行转置
            return w.transpose(0, 1) if self.fan_in_fan_out else w

        # 将模型设置为训练或者评估模式
        nn.Linear.train(self, mode)
        if mode:  # 训练模式
            if self.merge_weights and self.merged:
                # 解除权重的合并状态
                # Make sure that the weights are not merged
                if self.r > 0:
                    self.weight.data -= T(self.lora_B @ self.lora_A) * self.scaling
                self.merged = False
        else:  # 评估模式
            if self.merge_weights and not self.merged:
                # 合并权重
                # Merge the weights and mark it
                if self.r > 0:
                    self.weight.data += T(self.lora_B @ self.lora_A) * self.scaling
                self.merged = True

    def forward(self, x: torch.Tensor):
        def T(w):
            # 根据fan_in_fan_out的值，决定是否对权重矩阵进行转置
            return w.transpose(0, 1) if self.fan_in_fan_out else w

        if self.r > 0 and not self.merged:
            # 对输入x进行线性变换
            result = F.linear(x, T(self.weight), bias=self.bias)
            result += (self.lora_dropout(x) @ self.lora_A.transpose(0, 1) @ self.lora_B.transpose(0, 1)) * self.scaling
            return result
        else:
            # 直接使用F.linear函数对输入x进行线性变换，返回结果
            return F.linear(x, T(self.weight), bias=self.bias)


class MergedLinear(nn.Linear, LoRALayer):
    # LoRA implemented in a dense layer
    def __init__(
            self,
            in_features: int,
            out_features: int,
            r: int = 0,
            lora_alpha: int = 1,
            lora_dropout: float = 0.,
            enable_lora: List[bool] = [False],
            fan_in_fan_out: bool = False,
            merge_weights: bool = True,
            **kwargs
    ):
        nn.Linear.__init__(self, in_features, out_features, **kwargs)
        LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout,
                           merge_weights=merge_weights)
        assert out_features % len(enable_lora) == 0, \
            'The length of enable_lora must divide out_features'
        self.enable_lora = enable_lora
        self.fan_in_fan_out = fan_in_fan_out
        # Actual trainable parameters
        if r > 0 and any(enable_lora):
            self.lora_A = nn.Parameter(
                self.weight.new_zeros((r * sum(enable_lora), in_features)))
            self.lora_B = nn.Parameter(
                self.weight.new_zeros((out_features // len(enable_lora) * sum(enable_lora), r))
            )  # weights for Conv1D with groups=sum(enable_lora)
            self.scaling = self.lora_alpha / self.r
            # Freezing the pre-trained weight matrix
            self.weight.requires_grad = False
            # Compute the indices
            self.lora_ind = self.weight.new_zeros(
                (out_features,), dtype=torch.bool
            ).view(len(enable_lora), -1)
            self.lora_ind[enable_lora, :] = True
            self.lora_ind = self.lora_ind.view(-1)
        self.reset_parameters()
        if fan_in_fan_out:
            self.weight.data = self.weight.data.transpose(0, 1)

    def reset_parameters(self):
        nn.Linear.reset_parameters(self)
        if hasattr(self, 'lora_A'):
            # initialize A the same way as the default for nn.Linear and B to zero
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B)

    def zero_pad(self, x):
        result = x.new_zeros((len(self.lora_ind), *x.shape[1:]))
        result[self.lora_ind] = x
        return result

    def merge_AB(self):
        """用于合并lora_A和lora_B"""

        def T(w):
            # 是否对权重矩阵进行转置
            return w.transpose(0, 1) if self.fan_in_fan_out else w

        # 使用F.conv1d函数进行一维卷积操作
        delta_w = F.conv1d(
            self.lora_A.unsqueeze(0),  # 将lora_A的第一维度扩展为1
            self.lora_B.unsqueeze(-1),  # 将lora_B的最后一个维度扩展为1
            groups=sum(self.enable_lora)  # 设定卷积的组数
        ).squeeze(0)
        return T(self.zero_pad(delta_w))

    def train(self, mode: bool = True):
        def T(w):
            return w.transpose(0, 1) if self.fan_in_fan_out else w

        nn.Linear.train(self, mode)
        if mode:
            if self.merge_weights and self.merged:
                # Make sure that the weights are not merged
                if self.r > 0 and any(self.enable_lora):
                    self.weight.data -= self.merge_AB() * self.scaling
                self.merged = False
        else:
            if self.merge_weights and not self.merged:
                # Merge the weights and mark it
                if self.r > 0 and any(self.enable_lora):
                    self.weight.data += self.merge_AB() * self.scaling
                self.merged = True

    def forward(self, x: torch.Tensor):
        def T(w):
            return w.transpose(0, 1) if self.fan_in_fan_out else w

        if self.merged:
            return F.linear(x, T(self.weight), bias=self.bias)
        else:
            result = F.linear(x, T(self.weight), bias=self.bias)
            if self.r > 0:
                result += self.lora_dropout(x) @ T(self.merge_AB().T) * self.scaling
            return result


class ConvLoRA(nn.Module, LoRALayer):
    def __init__(self, conv_module, in_channels, out_channels, kernel_size, r=0, lora_alpha=1, lora_dropout=0.,
                 merge_weights=True, **kwargs):
        super(ConvLoRA, self).__init__()
        self.conv = conv_module(in_channels, out_channels, kernel_size, **kwargs)
        for name, param in self.conv.named_parameters():
            self.register_parameter(name, param)
        LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=merge_weights)
        assert isinstance(kernel_size, int)
        # Actual trainable parameters
        if r > 0:
            self.lora_A = nn.Parameter(
                self.conv.weight.new_zeros((r * kernel_size, in_channels * kernel_size))
            )
            self.lora_B = nn.Parameter(
                self.conv.weight.new_zeros((out_channels // self.conv.groups * kernel_size, r * kernel_size))
            )
            self.scaling = self.lora_alpha / self.r
            # Freezing the pre-trained weight matrix
            self.conv.weight.requires_grad = False
        self.reset_parameters()
        self.merged = False

    def reset_parameters(self):
        self.conv.reset_parameters()
        if hasattr(self, 'lora_A'):
            # initialize A the same way as the default for nn.Linear and B to zero
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B)

    def train(self, mode=True):
        super(ConvLoRA, self).train(mode)
        if mode:
            if self.merge_weights and self.merged:
                if self.r > 0:
                    # Make sure that the weights are not merged
                    self.conv.weight.data -= (self.lora_B @ self.lora_A).view(self.conv.weight.shape) * self.scaling
                self.merged = False
        else:
            if self.merge_weights and not self.merged:
                if self.r > 0:
                    # Merge the weights and mark it
                    self.conv.weight.data += (self.lora_B @ self.lora_A).view(self.conv.weight.shape) * self.scaling
                self.merged = True

    def forward(self, x):
        if self.r > 0 and not self.merged:
            return self.conv._conv_forward(
                x,
                self.conv.weight + (self.lora_B @ self.lora_A).view(self.conv.weight.shape) * self.scaling,
                self.conv.bias
            )
        return self.conv(x)


class Conv2d(ConvLoRA):
    def __init__(self, *args, **kwargs):
        super(Conv2d, self).__init__(nn.Conv2d, *args, **kwargs)


class Conv1d(ConvLoRA):
    def __init__(self, *args, **kwargs):
        super(Conv1d, self).__init__(nn.Conv1d, *args, **kwargs)


# Can Extend to other ones like this

class Conv3d(ConvLoRA):
    def __init__(self, *args, **kwargs):
        super(Conv3d, self).__init__(nn.Conv3d, *args, **kwargs)
